import torch
import torch.nn as nn
from torch.nn import init

class ChannelAttention(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.maxpool = nn.AdaptiveMaxPool2d(1)
        self.avgpool = nn.AdaptiveAvgPool2d(1)
        self.se = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, 1, bias=False),
            nn.GELU(),
            nn.Conv2d(channel // reduction, channel, 1, bias=False)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result = self.maxpool(x)
        avg_result = self.avgpool(x)
        max_out = self.se(max_result)
        avg_out = self.se(avg_result)
        out = self.sigmoid(max_out + avg_out)
        return out

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size=kernel_size, padding=kernel_size // 2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_result, _ = torch.max(x, dim=1, keepdim=True)
        avg_result = torch.mean(x, dim=1, keepdim=True)
        result = torch.cat([max_result, avg_result], dim=1)
        out = self.conv(result)
        out = self.sigmoid(out)
        return out

class CBAMBlock(nn.Module):
    def __init__(self, channel=512, reduction=16, kernel_size=7):
        super().__init__()
        self.ChannelAttention = ChannelAttention(channel=channel, reduction=reduction)
        self.SpatialAttention = SpatialAttention(kernel_size=kernel_size)

    def forward(self, x):
        # B, C, H, W = x.size()
        residual = x
        out = x * self.ChannelAttention(x)
        out = out * self.SpatialAttention(out)
        return residual + out

class CSA(nn.Module):
    def __init__(self, channel=512, reduction=16, kernel_size=7):
        super().__init__()
        self.ChannelAttention = ChannelAttention(channel=channel, reduction=reduction)
        self.SpatialAttention = SpatialAttention(kernel_size=kernel_size)

    def forward(self, x):
        # B, C, H, W = x.size()
        out = x * self.ChannelAttention(x)
        out = out * self.SpatialAttention(out)
        return out

class ConvNorm(nn.Module):
    def __init__(self,
                 in_planes,
                 out_planes,
                 kernel_size=1,
                 stride=1,
                 padding=0,
                 dilation=1,
                 groups=1,
                 with_bn=True,
                 norm_type='BN'):
        super().__init__()
        self.with_bn = with_bn
        self.conv = nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding, dilation, groups)
        if self.with_bn:
            self.norm = nn.BatchNorm2d(out_planes)
            torch.nn.init.constant_(self.norm.weight, 1)
            torch.nn.init.constant_(self.norm.bias, 0)

    def forward(self, x):
        #(B, C, H, W)
        x = self.conv(x)
        if self.with_bn:
            x = self.norm(x)
        return x


class BF3M(nn.Module):
    def __init__(self, dim, reduction=8, kernel_size=3):
        super(BF3M, self).__init__()
        self.csa = CSA(dim, reduction, kernel_size)
        self.dwconv1 = ConvNorm(dim, dim, kernel_size, 1, (kernel_size - 1) // 2, groups=dim, with_bn=True)
        self.dwconv2 = ConvNorm(dim, dim, kernel_size, 1, (kernel_size - 1) // 2, groups=dim, with_bn=True)
        self.pwconv = ConvNorm(dim * 2, dim, 1, with_bn=True)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x, y):
        input = x + y
        residual_x = x
        residual_y = y
        attn = self.sigmoid(self.csa(input))
        x = self.dwconv1(x) * attn + residual_x
        y = self.dwconv2(y) * (1 - attn) + residual_y
        output = torch.cat((x, y), dim=1)
        # output = self.pwconv(output)
        return output


if __name__ == '__main__':
    x = torch.rand(1, 192, 56, 56)
    y = torch.rand(1, 192, 56, 56)
    model = BF3M(dim=192, reduction=8, kernel_size=3)
    print(model(x, y).shape) # torch.Size([1, 384, 56, 56])